# Copyright (c) OpenMMLab. All rights reserved.
import logging
import sys
from contextlib import contextmanager
from logging import Logger, LogRecord
from typing import Optional


from transformers import PretrainedConfig

logger_initialized = {}


class _ASNI_COLOR:
    BRIGHT_RED = '\033[91m'
    RED = '\033[31m'
    YELLOW = '\033[33m'
    WHITE = '\033[37m'
    GREEN = '\033[32m'


class ColorFormatter(logging.Formatter):

    _LEVELNAME_COLOR_MAP = dict(CRITICAL=_ASNI_COLOR.BRIGHT_RED,
                                ERROR=_ASNI_COLOR.RED,
                                WARN=_ASNI_COLOR.YELLOW,
                                WARNING=_ASNI_COLOR.YELLOW,
                                INFO=_ASNI_COLOR.WHITE,
                                DEBUG=_ASNI_COLOR.GREEN)

    _RESET_COLOR = '\033[0m'

    def format(self, record: LogRecord):
        """format."""
        if sys.platform == 'win32':
            # windows does not support ASNI color
            return super().format(record)
        levelname = record.levelname
        level_color = self._LEVELNAME_COLOR_MAP.get(levelname,
                                                    self._RESET_COLOR)
        levelname = f'{level_color}{levelname}{self._RESET_COLOR}'
        record.levelname = levelname
        return super().format(record)


class FilterDuplicateWarning(logging.Filter):
    """Filter the repeated warning message.

    Args:
        name (str): name of the filter.
    """

    def __init__(self, name: str = 'lmdeploy'):
        super().__init__(name)
        self.seen: set = set()

    def filter(self, record: LogRecord) -> bool:
        """Filter the repeated warning message.

        Args:
            record (LogRecord): The log record.

        Returns:
            bool: Whether to output the log record.
        """
        if record.levelno != logging.WARNING:
            return True

        if record.msg not in self.seen:
            self.seen.add(record.msg)
            return True
        return False

# add by haoyong_liu@163.com 2024-11-18 start
_FORMAT = '%(asctime)s %(levelname)s %(filename)s:%(lineno)d - %(message)s'
# remove by haoyong_liu@163.com 2024-11-18 start
'''
_FORMAT = '%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d' \
          ' - %(message)s'
'''


def get_logger(name: Optional[str] = None,
               log_file: Optional[str] = None,
               log_level: int = logging.INFO,
               file_mode: str = 'w',
               log_formatter: str = _FORMAT) -> Logger:
    """Initialize and get a logger by name.

    If the logger has not been initialized, this method will initialize the
    logger by adding one or two handlers, otherwise the initialized logger will
    be directly returned. During initialization, a StreamHandler will always be
    added. If `log_file` is specified, a FileHandler will also be added.
    Args:
        name (str): Logger name.
        log_file (str | None): The log filename. If specified, a FileHandler
            will be added to the logger.
        log_level (int): The logger level.
        file_mode (str): The file mode used in opening log file.
            Defaults to 'w'.
        log_formatter (str): The logger output format.
    Returns:
        logging.Logger: The expected logger.
    """
    logger = logging.getLogger(name)
    if name in logger_initialized:
        return logger
    # handle hierarchical names
    # e.g., logger "a" is initialized, then logger "a.b" will skip the
    # initialization since it is a child of "a".
    for logger_name in logger_initialized:
        if name.startswith(logger_name):
            return logger

    # handle duplicate logs to the console
    for handler in logger.root.handlers:
        if type(handler) is logging.StreamHandler:
            handler.setLevel(logging.ERROR)

    stream_handler = logging.StreamHandler(stream=sys.stdout)
    
    # add by haoyong_liu@163.com 2024-11-18 start
    #formatter = ColorFormatter(log_formatter)
    formatter = logging.Formatter(log_formatter)
    stream_handler.setFormatter(formatter)
    stream_handler.setLevel(logging.DEBUG)
    stream_handler.addFilter(FilterDuplicateWarning(name))
    handlers = [stream_handler]
    logger.addHandler(stream_handler)
    if log_file is None:
        import os
        if os.getenv('TM_LOG_FILE') is not None:
            log_file = os.environ['TM_LOG_FILE']
        else:
            log_file = name + '.log'
    if log_file is not None:
        from logging.handlers import TimedRotatingFileHandler
        # 创建一个handler，用于写入日志文件，按天分割
        file_handler = TimedRotatingFileHandler(
            filename=log_file,  # 日志文件路径
            when='midnight',  # 在午夜进行日志切割
            interval=1,  # 间隔为1天
            backupCount=120,  # 保留最近120天的日志文件
            encoding='utf-8'  # 设置日志编码
        )
        file_handler.suffix = "%Y-%m-%d"
        # 创建一个formatter，定义日志的输出格式
        file_formatter = logging.Formatter(_FORMAT)
        file_handler.setFormatter(file_formatter)
        file_handler.setLevel(log_level)
        file_handler.addFilter(FilterDuplicateWarning(name))
        # 给logger添加handler
        handlers.append(file_handler)
        logger.addHandler(file_handler)
    # add by haoyong_liu@163.com 2024-11-18	end
    # remove by haoyong_liu@163.com 2024-11-18 start
    '''
    handlers = [stream_handler]
    if log_file is not None:
        # Here, the default behaviour of the official logger is 'a'. Thus, we
        # provide an interface to change the file mode to the default
        # behaviour.
        file_handler = logging.FileHandler(log_file, file_mode)
        handlers.append(file_handler)  

    formatter = ColorFormatter(log_formatter)
    for handler in handlers:
        handler.setFormatter(formatter)
        handler.setLevel(logging.DEBUG)
        handler.addFilter(FilterDuplicateWarning(name))
        logger.addHandler(handler)
    '''
    # remove by haoyong_liu@163.com 2024-11-18 end
    logger.setLevel(log_level)
    logger.propagate = False
    logger_initialized[name] = True

    return logger
